package cc.autoapi.pucong.auto.flow.core.flow.execNode.ai.send.handler;


import cc.autoapi.pucong.auto.flow.core.flow.execNode.ai.send.AbstractAiSendSendTool;
import cc.autoapi.pucong.auto.flow.core.flow.execNode.exception.AutoDesignException;
import cc.autoapi.pucong.auto.flow.core.flow.execNode.streamBus.IStreamSubscriber;
import cc.autoapi.pucong.auto.flow.core.flow.execNode.streamBus.StreamPublisher;
import cc.autoapi.pucong.auto.flow.core.flow.execNode.utils.FlowThreadPoolUtils;
import cn.hutool.http.HttpRequest;
import cn.hutool.http.HttpResponse;
import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.PrintWriter;
import java.net.HttpURLConnection;
import java.net.URL;
import java.time.LocalDateTime;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

/**
 * @description:
 * @author: cong.pu
 * @create: 2024-10-14 09:01
 **/
@Slf4j
public class TransAiHandlerSend extends AbstractAiSendSendTool {

    private String appId;
    private String appKey;
    private String token;

    private LocalDateTime authTime;
    private Integer expiresIn;

    public TransAiHandlerSend(String appId, String appKey) {
        this.appId = appId;
        this.appKey = appKey;
    }

    public String getToken() {
        if (StringUtils.isBlank(token)) {
            initToken();
        } else {
            // 是否过期
            LocalDateTime now = LocalDateTime.now();
            LocalDateTime authTime2 = authTime.plusSeconds(expiresIn);
            if (authTime2.isBefore(now)) {
                initToken();
            }
        }
        return token;
    }

    @Override
    public void initToken() {
        String url = "https://pfgateway.transsion.com:9199/service-pf-open-gateway/oauth/token";
        String bodyParams = "{}";
        Map<String, Object> query = new HashMap<>(16);
        Map<String, String> header = new HashMap<>(16);
        query.put("grant_type", "client_credentials");
        query.put("client_id", appId);
        query.put("client_secret", appKey);

        String urlWithParam = url;
        Set<Map.Entry<String, Object>> entries = query.entrySet();
        if (entries.size() > 0) {
            urlWithParam = url + "?";
            for (Map.Entry<String, Object> entry : entries) {
                urlWithParam += entry.getKey() + "=" + entry.getValue() + "&";
            }
            if (!query.isEmpty()) {
                urlWithParam = urlWithParam.substring(0, urlWithParam.length() - 1);
            }
        }
        HttpResponse response = HttpRequest.post(urlWithParam)
                .addHeaders(header)
                .body(bodyParams)
                .contentType("application/json")
                .execute();
        int status = response.getStatus();
        String res = response.body();
        if (status == 200) {
            JSONObject jsonObject = JSONObject.parseObject(res);
            JSONObject data = jsonObject.getJSONObject("data");
            token = data.getString("access_token");
            expiresIn = data.getInteger("expires_in");
            authTime = LocalDateTime.now();
        } else {
            throw new AutoDesignException(res);
        }
    }

    @Override
    public String sendToAi(String question) {
        String url = "https://gptgateway.transsion.com/api/azure/openai/chatCompletion";

        Map<String, Object> queryMap = new HashMap<>(16);
        queryMap.put("deploymentId", "gpt-4o");
        Map<String, String> headerMap = new HashMap<>(16);
        headerMap.put("Authorization", "Bearer " + getToken());

        Map<String, Object> bodyMap = new HashMap<>(16);
        JSONArray messages = new JSONArray();
        String system = getSystem();
        if (StringUtils.isNotBlank(system)) {
            JSONObject systemMsg = new JSONObject();
            systemMsg.put("role", "system");
            systemMsg.put("content", system);
            messages.add(systemMsg);
        }
        JSONObject userMsg = new JSONObject();
        userMsg.put("role", "user");
        userMsg.put("content", question);
        messages.add(userMsg);
        bodyMap.put("messages", messages);
        bodyMap.put("temperature", 0.2);


        String urlWithParam = url;
        Set<Map.Entry<String, Object>> entries = queryMap.entrySet();
        if (entries.size() > 0) {
            urlWithParam = url + "?";
            for (Map.Entry<String, Object> entry : entries) {
                urlWithParam += entry.getKey() + "=" + entry.getValue() + "&";
            }
            if (!queryMap.isEmpty()) {
                urlWithParam = urlWithParam.substring(0, urlWithParam.length() - 1);
            }
        }
        String bodyParams = JSONObject.toJSONString(bodyMap);
        HttpResponse response = HttpRequest.post(urlWithParam)
                .addHeaders(headerMap)
                .body(bodyParams)
                .contentType("application/json;charset=utf-8")
                .execute();
        int status = response.getStatus();
        String res = response.body();
        if (status == 200) {
            JSONObject jsonObject = JSONObject.parseObject(res);
            String msg = jsonObject.getJSONObject("data").getJSONArray("choices").getJSONObject(0).getJSONObject("message").getString("content");
            return msg;
        }
        throw new AutoDesignException(res);
    }

    @Override
    public JSONObject sendToolsToAi(JSONArray msgArrays) {
        String url = "https://gptgateway.transsion.com/api/azure/openai/chatCompletion";

        Map<String, Object> queryMap = new HashMap<>(16);
        queryMap.put("deploymentId", "gpt-4o");
        Map<String, String> headerMap = new HashMap<>(16);
        headerMap.put("Authorization", "Bearer " + getToken());
        Map<String, Object> bodyMap = new HashMap<>(16);
        bodyMap.put("messages", msgArrays);
        bodyMap.put("temperature", 0.2);
        bodyMap.put("toolChoice", "auto");

        if (getTools() != null) {
            bodyMap.put("tools", getTools());
        }
        String urlWithParam = url;
        Set<Map.Entry<String, Object>> entries = queryMap.entrySet();
        if (entries.size() > 0) {
            urlWithParam = url + "?";
            for (Map.Entry<String, Object> entry : entries) {
                urlWithParam += entry.getKey() + "=" + entry.getValue() + "&";
            }
            if (!queryMap.isEmpty()) {
                urlWithParam = urlWithParam.substring(0, urlWithParam.length() - 1);
            }
        }
        String bodyParams = JSONObject.toJSONString(bodyMap);
        HttpResponse response = HttpRequest.post(urlWithParam)
                .addHeaders(headerMap)
                .body(bodyParams)
                .contentType("application/json;charset=utf-8")
                .execute();
        int status = response.getStatus();
        String res = response.body();
        if (status == 200) {
            JSONObject jsonObject = JSONObject.parseObject(res);
            String resStatus = jsonObject.getString("code");
            if (!"200".equals(resStatus)) {
                throw new AutoDesignException(jsonObject.getString("message"));
            }
            JSONObject jsonObject1 = jsonObject.getJSONObject("data").getJSONArray("choices").getJSONObject(0).getJSONObject("message");
            return jsonObject1;
        }
        throw new AutoDesignException(res);
    }

    @Override
    public IStreamSubscriber sendToAiStream(String question) {
        String urlPath = "https://gptgateway.transsion.com/api/azure/openai/chatCompletion";

        Map<String, Object> queryMap = new HashMap<>(16);
        queryMap.put("deploymentId", "gpt-4o");
        Map<String, String> headerMap = new HashMap<>(16);
        headerMap.put("Authorization", "Bearer " + token);

        Map<String, Object> bodyMap = new HashMap<>(16);
        JSONArray messages = new JSONArray();
        JSONObject systemMsg = new JSONObject();
        systemMsg.put("role", "system");
        systemMsg.put("content", getSystem());
        JSONObject userMsg = new JSONObject();
        userMsg.put("role", "user");
        userMsg.put("content", question);
        messages.add(systemMsg);
        messages.add(userMsg);
        bodyMap.put("messages", messages);
        bodyMap.put("temperature", 0.2);


        String urlWithParam = urlPath;
        Set<Map.Entry<String, Object>> entries = queryMap.entrySet();
        if (entries.size() > 0) {
            urlWithParam = urlPath + "?";
            for (Map.Entry<String, Object> entry : entries) {
                urlWithParam += entry.getKey() + "=" + entry.getValue() + "&";
            }
            if (!queryMap.isEmpty()) {
                urlWithParam = urlWithParam.substring(0, urlWithParam.length() - 1);
            }
        }
        String bodyParams = JSONObject.toJSONString(bodyMap);

        StreamPublisher publisher = new StreamPublisher();
        String finalUrlWithParam = urlWithParam;
        FlowThreadPoolUtils.executor(() -> {
            PrintWriter out = null;
            BufferedReader in = null;
            try {
                URL url = new URL(finalUrlWithParam);
                HttpURLConnection conn = (HttpURLConnection) url.openConnection();
                //需要用到输出流
                conn.setDoOutput(true);
                //需要用到输入流
                conn.setDoInput(true);
                conn.setRequestMethod("POST");
                conn.setRequestProperty("Accept", "*/*");
                conn.setRequestProperty("Connection", "keep-alive");
                conn.setRequestProperty("content-type", "application/json");
                // 使用输出流添加数据至请求体
                if (StringUtils.isNotBlank(bodyParams)) {
                    // 获取URLConnection对象对应的输出流
                    out = new PrintWriter(conn.getOutputStream());
                    // 发送请求参数
                    out.print(bodyParams);
                    // flush输出流的缓冲
                    out.flush();
                }
                conn.connect();
                // 定义BufferedReader输入流来读取URL的响应
                in = new BufferedReader(new InputStreamReader(conn.getInputStream()));
                String line = null;
                StringBuilder sbf = new StringBuilder();
                while ((line = in.readLine()) != null) {
                    line = line.replaceAll("data:", "").trim();
                    if (StringUtils.isBlank(line) && sbf.length() > 0) {
                        String string = sbf.toString();
                        publisher.send(string);
                        sbf = new StringBuilder();
                    }
                    sbf.append(line);
                }
                if (sbf.length() > 0) {
                    publisher.send(sbf.toString());
                }
            } catch (Exception e) {
                log.error("发送流式请求出现异常！", e);
            } finally {
                try {
                    if (out != null) {
                        out.close();
                    }
                    if (in != null) {
                        in.close();
                    }
                } catch (IOException ex) {
                    ex.printStackTrace();
                }
                publisher.complete();
            }
        });
        return publisher.getSubscriber();
    }
}
